package org.jmeter.tcp.helmet.client;

import io.netty.bootstrap.Bootstrap;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.nio.NioSocketChannel;
import org.jmeter.tcp.helmet.client.NettyDecoder;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.channels.SocketChannel;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Supplier;

/**
 * @author liangsu
 * @version v1.0
 * @Description
 * @Date 2020/11/17 11:21
 * @since
 */
public class ConnectionContext {
    private static EventLoopGroup group = new NioEventLoopGroup(8);

    static ThreadLocal<Map<String, Object>> socketChannelMap = new ThreadLocal<Map<String, Object>>(){
        @Override
        protected Map<String, Object> initialValue() {
            return new HashMap<String, Object>();
        }
    };

    public static Object getClient(String ip, Integer port, String devId, Supplier supplier) throws IOException {
        String key = ip + port + devId;
        Map<String, Object> map = socketChannelMap.get();
        Object client = map.get(key);
        if(client == null){
            client = supplier.get();
            map.put(key, client);
        }

        return client;
    }

    public static SocketChannel getConnection(String ip, Integer port, String devId) throws IOException {
        SocketChannel socketChannel = (SocketChannel)getClient(ip, port, devId, () -> {
            SocketChannel client = null;
            try {
                client = SocketChannel.open();
                client.connect(new InetSocketAddress(ip, port));
            } catch (IOException e) {
                e.printStackTrace();
            }

            return client;
        });

        return socketChannel;
    }


    public static NioSocketChannel getNettyClient(String ip, Integer port, String devId) throws IOException {
        NioSocketChannel channel = (NioSocketChannel)getClient(ip, port, devId, () -> {

            Bootstrap bootstrap = new Bootstrap();
            ChannelFuture connect = bootstrap
                    .group(group)
                    .channel(NioSocketChannel.class)
                    .option(ChannelOption.SO_KEEPALIVE, true)
                    .handler(new ChannelInitializer<NioSocketChannel>() {
                        @Override
                        protected void initChannel(NioSocketChannel ch) throws Exception {
                            ChannelPipeline pipeline = ch.pipeline();
                            pipeline.addLast(new NettyDecoder());
                        }
                    })
                    .connect(new InetSocketAddress(ip, port));

            NioSocketChannel client = null;
            try {
                client = (NioSocketChannel) connect.sync().channel();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }

            return client;
        });

        return channel;
    }

}
